package com.ly.algorithm.DataStructures.huffmantree;

import java.util.*;

/**
 * @Classname HuffManCode
 * @Description TODO
 * @Date 2020/5/2 19:57
 * @Author 冷心影翼
 */
public class HuffManCode {

    public Map<Byte,String> huffManCodeTable = new HashMap<>();

    public CodeNode root;

    public String huffManCode;

    public byte[] huffManCodeBytes;

    public void createHuffManTable(CodeNode node,String code,StringBuilder sb) {
        StringBuilder strBuilder = new StringBuilder(sb);
        strBuilder.append(code);

        if (node!=null) {
            if(node.data==null) {
                createHuffManTable(node.left,"0",strBuilder);
                createHuffManTable(node.right,"1",strBuilder);
            }else {
                huffManCodeTable.put(node.data,strBuilder.toString());
            }
        }
    }

    private void createHuffManCode(List<CodeNode> nodes) {
        while (nodes.size()>1) {
            nodes.sort(Comparator.comparingInt( n -> n.weight));
            CodeNode left = nodes.get(0);
            CodeNode right = nodes.get(1);
            CodeNode parent = new CodeNode(null,left.weight+right.weight);
            parent.left = left;
            parent.right = right;
            nodes.remove(left);
            nodes.remove(right);
            nodes.add(parent);
        }
        root = nodes.get(0);
    }

    public HuffManCode(byte[] bytes) {
        List<CodeNode> codeNodes = getNodes(bytes);
        createHuffManCode(codeNodes);
        createHuffManTable(root,"",new StringBuilder());
    }

    private List<CodeNode> getNodes(byte[] bytes) {
        List<CodeNode> nodes = new ArrayList<>();
        Map<Byte,Integer> map = new HashMap<>();
        for (byte b : bytes) {
            Integer integer = map.get(b);
            if(integer == null) {
                map.put(b,1);
            }else {
                map.put(b,integer+1);
            }
        }

        for(Map.Entry<Byte, Integer> entry : map.entrySet()) {
            nodes.add(new CodeNode(entry.getKey(),entry.getValue()));
        }

        return nodes;
    }

    public void preOrder() {
        if(root!=null) {
            root.preOrder();
        }
    }

    public void zip(byte[] bytes) {
        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes) {
            String s = huffManCodeTable.get(b);
            stringBuilder.append(s);
        }
        huffManCode = stringBuilder.toString();

        int length = (huffManCode.length()+7)/8;
        byte[] bytes2 = new byte[length];

        int index = 0;
        for(int i=0;i<huffManCode.length();i+=8) {
            byte b;
            if(i+8>huffManCode.length()) {
                b = (byte)Integer.parseInt(stringBuilder.substring(i),2);
            }else {
                b = (byte)Integer.parseInt(stringBuilder.substring(i,i+8),2);
            }
            bytes2[index++] = b;
        }
        huffManCodeBytes = bytes2;
    }



    private String byteToString(boolean flag,byte b) {
        int temp = b;
        if(flag||temp<0) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if(flag) {
            return str.substring(str.length()-8);
        }else {
            return str;
        }
    }

    public String decode(byte[] bytes) {
        StringBuilder stringBuilder = new StringBuilder();
        for(int i=0;i<bytes.length;i++) {
            if(i== bytes.length-1) {
                stringBuilder.append(byteToString(false,bytes[i]));
            }else {
                stringBuilder.append(byteToString(true,bytes[i]));
            }
        }

        System.out.println(stringBuilder.toString());
        int index = 0;
        List<Byte> list = new ArrayList<>();
        Map<String,Byte> map = new HashMap<>();
        for(Map.Entry<Byte,String> entry : huffManCodeTable.entrySet()) {
            map.put(entry.getValue(),entry.getKey());
        }
        int count = 1;
        while (true) {
            if(index + count >= stringBuilder.length()) {
                String str = stringBuilder.substring(index);
                list.add(map.get(str));
                break;
            }
            String str = stringBuilder.substring(index,index+count);
            if(map.containsKey(str)) {
                list.add(map.get(str));
                index += count;
                count = 1;
            }else {
                count++;
            }
        }

        byte[] res = new byte[list.size()];
        for(int i=0;i<list.size();i++) {
            res[i] = list.get(i);
        }
        return new String(res);
    }


}


class CodeNode {
    int weight;
    Byte data;
    CodeNode left;
    CodeNode right;

    public CodeNode(Byte data,int weight) {
        this.data = data;
        this.weight = weight;
    }

    public void preOrder() {
        System.out.println(this);
        if(left!=null) {
            left.preOrder();
        }
        if(right!=null) {
            right.preOrder();
        }
    }

    @Override
    public String toString() {
        return "Node" +
                "weight=" + weight +
                ", data=" + data ;
    }


}


class Test2{
    public static void main(String[] args) {
//        String str = "i like like like 13s qewe qu beq 1 wo p ll java do you like a javab";
        String str = "i like like like java do you like a java";
        byte[] bytes = str.getBytes();
        HuffManCode huffManCode = new HuffManCode(bytes);
//        huffManCode.preOrder();
//        System.out.println(huffManCode.huffManCodeTable);
        huffManCode.zip(bytes);
        System.out.println(huffManCode.huffManCode);
//        System.out.println(huffManCode.huffManCode.length());
//        System.out.println(Arrays.toString(huffManCode.huffManCodeBytes));
//        System.out.println(huffManCode.huffManCodeBytes.length);

        System.out.println(huffManCode.decode(huffManCode.huffManCodeBytes));
    }
}
